'''
description: a async port proxy
features:
    1. 支持多转发规则
    2. 支持持久化转发规则
    3. 支持文件/命令行参数设置转发规则
    4. 支持运行时动态添加/删除转发规则
todos:
    1. 支持统计功能：会话，连接数。
@author: zhouchangxun@outlook.com
'''
import asyncio
import sys
import uuid
import json


class PortProxy:
    def __init__(self):
        # 存储转发规则的字典，键为监听端口，值为 (server, remote_addr, remote_port) 元组
        self.forward_rules = {}
        self.CONFIG_FILE = 'forward_rules.json'

    # forwarder: server <-> client
    async def forwarder(self, reader, writer):
        while not reader.at_eof():
            data = await reader.read(256)
            writer.write(data)
            await writer.drain()
        writer.close()

    # new connection handler
    async def handle(self, cli_reader, cli_writer):
        sid = 'sid-' + str(uuid.uuid4()).split('-')[-1]
        addr = cli_writer.get_extra_info('peername')
        print(f'> [{sid}] Receive connection: {addr}!')
        listen_port = cli_writer.get_extra_info('sockname')[1]
        # 找到对应的转发规则
        rule = self.forward_rules.get(listen_port)
        if rule is None:
            # 没有找到对应的规则，关闭连接
            cli_writer.close()
            return
        r_addr, r_port = rule[1:]
        # connect to remote
        r_reader, r_writer = await asyncio.open_connection(r_addr, r_port)
        print(f'> [{sid}] Connect to remote: {r_addr}:{r_port}')
        ret = await asyncio.gather(
            self.forwarder(cli_reader, r_writer),
            self.forwarder(r_reader, cli_writer)
        )
        print(f'> [{sid}] Close connection: {addr} !')

    async def add_forward_rule_async(self, listen_port, remote_addr, remote_port):
        """
        异步添加转发规则
        :param listen_port: 监听端口
        :param remote_addr: 远程地址
        :param remote_port: 远程端口
        """
        if listen_port in self.forward_rules:
            print(f'Forward rule for listen port {listen_port} already exists.')
            return
        server = await asyncio.start_server(self.handle, "0.0.0.0", listen_port)
        addr = server.sockets[0].getsockname()
        print(f'Serving on {addr} for {remote_addr}:{remote_port}')
        # 将服务器添加到字典中
        self.forward_rules[listen_port] = (server, remote_addr, remote_port)
        # 创建一个任务，用于运行新的监听服务
        asyncio.create_task(server.serve_forever())
        print(f'Added forward rule: ({listen_port}, {remote_addr}, {remote_port})')
        # 持久化到文件
        self._save_config()

    def remove_forward_rule(self, listen_port):
        """
        删除转发规则
        :param listen_port: 监听端口
        """
        rule = self.forward_rules.get(listen_port)
        if rule is None:
            print(f'No forward rule found for listen port {listen_port}')
            return
        server, r_addr, r_port = rule
        server.close()
        del self.forward_rules[listen_port]
        print(f'Removed forward rule: ({listen_port}, {r_addr}, {r_port})')
        # 持久化到文件
        self._save_config()

    def _load_config(self):
        try:
            with open(self.CONFIG_FILE, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            return {}

    def _save_config(self):
        config = {int(port): (r_addr, r_port) for port, (_, r_addr, r_port) in self.forward_rules.items()}
        with open(self.CONFIG_FILE, 'w') as f:
            json.dump(config, f, indent=4)

    def stop(self):
        for listen_port, (server, _, _) in self.forward_rules.items():
            print(f"closing server *:{listen_port}")
            server.close()
    async def main(self):
        # 1. read config from file
        print(f'Loading file config...')
        startup_rules = self._load_config()
        # 2. read config from cmdline if exists
        print(f'Loading cmdline config...')
        try:
            if len(sys.argv) > 3:
                port = int(sys.argv[1])
                r_addr = sys.argv[2]
                r_port = int(sys.argv[3])
                startup_rules[str(port)] = [r_addr, r_port]
        except Exception:
            print(f'> usage:\r\n  python {sys.argv[0]} [listen_port] [romote_addr] [remote_port]')
        if not startup_rules:
            print(f'Exited: No forward rule found in config file or cmdline.')
            return
        # 3. add startup rules
        print(f'rules: {startup_rules}')
        for rule in startup_rules.items():
            listen_port, (r_addr, r_port) = rule
            await self.add_forward_rule_async(listen_port, r_addr, r_port)
        # 4. keep main() running.
        while True:
            await asyncio.sleep(1)


if __name__ == "__main__":
    proxy = PortProxy()
    try:
        asyncio.run(proxy.main())
    except KeyboardInterrupt:
        print('Ctrl+C pressed. Shutting down...')
        # 关闭所有监听
        proxy.stop()
        print('> All servers have been shut down. bye!')